#include "solver.h"
#include <cmath>
#include <iostream>
#include <algorithm>
#include <fstream>

const double g = 9.81;  // 重力加速度 (m/s^2)
const double PI = M_PI; // 圆周率

MultiphaseFlowSolver::MultiphaseFlowSolver(const Fluid& fluid, const PipeParameters& params)
    : fluid_(fluid), params_(params) {
    // 初始化网格参数
    numGrids_ = 100;  // 将管道分成100个网格
    dx_ = params.length / (numGrids_ - 1);
    dt_ = params.minTimeStep;  // 使用最小时间步长作为初始步长
    
    // 初始化所有向量大小
    alpha_g_.resize(numGrids_);
    alpha_d_.resize(numGrids_);
    alpha_l_.resize(numGrids_);
    V_g_.resize(numGrids_);
    V_d_.resize(numGrids_);
    V_l_.resize(numGrids_);
    pressure_.resize(numGrids_);
    temperature_.resize(numGrids_);
    
    // 初始化物性数组
    gasDensity_.resize(numGrids_);
    oilDensity_.resize(numGrids_);
    waterDensity_.resize(numGrids_);
    gasViscosity_.resize(numGrids_);
    oilViscosity_.resize(numGrids_);
    waterViscosity_.resize(numGrids_);
    
    // 初始化通量存储
    gas_flux_old_.resize(numGrids_);
    liquid_flux_old_.resize(numGrids_);
    
    // 初始化压力修正量
    pressure_correction_.resize(numGrids_);
    
    // 初始化能量通量数组
    energy_flux_old_.resize(numGrids_);
    
    // 调用初始化流场函数
    initializeFlow();
}

MultiphaseFlowSolver::Results MultiphaseFlowSolver::solve() {
    Results results;
    
    // 初始化结果数组
    results.pressure.resize(numGrids_);
    results.temperature.resize(numGrids_);
    results.gasHoldup.resize(numGrids_);
    results.oilHoldup.resize(numGrids_);
    results.waterHoldup.resize(numGrids_);
    results.gasVelocity.resize(numGrids_);
    results.oilVelocity.resize(numGrids_);
    results.waterVelocity.resize(numGrids_);
    
    // 初始化流场
    initializeFlow();
    
    double time = 0.0;
    const double max_time = params_.simulationTime; // 只模拟60秒，足够达到稳态
    
    // 打印初始条件
    std::cout << "\nInitial conditions:" << std::endl;
    for (int i = 0; i < numGrids_; i += numGrids_/10) {
        std::cout << "Grid " << i << ": "
                  << "P=" << pressure_[i] << " Pa, "
                  << "T=" << temperature_[i] << " K, "
                  << "alpha_g=" << alpha_g_[i] << ", "
                  << "rho_g=" << gasDensity_[i] << " kg/m3" << std::endl;
    }
    
    // 时间推进循环
    while (time < max_time) {
        // 1. 求解动量方程
        solveGasDropletMomentum();
        
        // 2. 求解压力方程
        solvePressureCorrection();
        
        // 3. 更新相分数
        solveMassConservation();
        
        // 4. 求解能量方程
        solveEnergyEquation();
        
        // 5. 更新物性
        updateProperties();
        
        // 检查数值不稳定性
        bool stable = true;
        for (int i = 0; i < numGrids_; i++) {
            if (!std::isfinite(pressure_[i]) || !std::isfinite(V_g_[i]) || 
                !std::isfinite(V_l_[i]) || !std::isfinite(temperature_[i])) {
                std::cout << "Warning: Numerical instability at time = " 
                          << time << ", grid = " << i << std::endl;
                stable = false;
                break;
            }
        }
        
        if (!stable) break;
        
        time += dt_;
    }
    
    // 保存最终结果
    saveResults(results, time);
    std::cout << "Solution completed at time = " << time << " seconds." << std::endl;
    
    return results;
}

void MultiphaseFlowSolver::initializeFlow() {
    // 线性初始化压力
    for (int i = 0; i < numGrids_; i++) {
        pressure_[i] = params_.inletPressure - 
                      (params_.inletPressure - params_.outletPressure) * i / (numGrids_ - 1);
    }
    
    // 初始温度为入口温度
    for (int i = 0; i < numGrids_; i++) {
        temperature_[i] = params_.inletTemp;
    }
    
    // 初始相分数
    for (int i = 0; i < numGrids_; i++) {
        alpha_g_[i] = params_.inletGasVolumeFraction;
        alpha_l_[i] = 1.0 - alpha_g_[i];
        alpha_d_[i] = 0.0;
    }
    
    // 初始速度
    for (int i = 0; i < numGrids_; i++) {
        V_g_[i] = params_.inletGasVelocity;
        V_l_[i] = params_.inletLiquidVelocity;
        V_d_[i] = V_l_[i];
    }
    
    // 更新物性
    updateProperties();
}

void MultiphaseFlowSolver::solveGasDropletMomentum() {
    for (int i = 1; i < numGrids_ - 1; i++) {
        // 保存旧的速度
        double old_Vg = V_g_[i];
        double old_Vl = V_l_[i];
        
        // 计算压力梯度
        double dp_dx = (pressure_[i+1] - pressure_[i-1]) / (2.0 * dx_);
        
        // 计算摩擦系数
        double f_g = calculateFrictionFactor(i, "gas");
        double f_l = calculateFrictionFactor(i, "liquid");
        
        // 壁面摩擦阻力
        double tau_g = 0.5 * f_g * gasDensity_[i] * V_g_[i] * std::abs(V_g_[i]);
        double tau_l = 0.5 * f_l * oilDensity_[i] * V_l_[i] * std::abs(V_l_[i]);
        
        // 相间力 - 简化为拖曳力
        double k_i = 0.44;  // 拖曳系数
        double F_drag = k_i * gasDensity_[i] * std::abs(V_g_[i] - V_l_[i]) * (V_g_[i] - V_l_[i]);
        
        // 动量方程 - 气相
        double S_g = -1.0/gasDensity_[i] * dp_dx 
                   - 4.0*tau_g/(gasDensity_[i] * params_.diameter) 
                   - F_drag/(alpha_g_[i] * gasDensity_[i])
                   - g * std::sin(params_.inclination);
        
        // 动量方程 - 液相
        double S_l = -1.0/oilDensity_[i] * dp_dx 
                   - 4.0*tau_l/(oilDensity_[i] * params_.diameter) 
                   + F_drag/(alpha_l_[i] * oilDensity_[i])
                   - g * std::sin(params_.inclination);
        
        // 使用更小的时间步长和松弛
        double local_dt = std::min(dt_, 0.001);
        double omega = 0.2;
        
        // 更新速度
        V_g_[i] = (1.0 - omega) * old_Vg + omega * (old_Vg + local_dt * S_g);
        V_l_[i] = (1.0 - omega) * old_Vl + omega * (old_Vl + local_dt * S_l);
        
        // 限制速度范围
        V_g_[i] = std::max(0.1, std::min(10.0, V_g_[i]));
        V_l_[i] = std::max(0.1, std::min(5.0, V_l_[i]));
        V_d_[i] = V_l_[i];  // 液滴速度等于液相速度
    }
    
    // 应用边界条件
    V_g_[0] = params_.inletGasVelocity;
    V_l_[0] = params_.inletLiquidVelocity;
    V_d_[0] = V_l_[0];
    
    V_g_[numGrids_-1] = V_g_[numGrids_-2];
    V_l_[numGrids_-1] = V_l_[numGrids_-2];
    V_d_[numGrids_-1] = V_d_[numGrids_-2];
}

void MultiphaseFlowSolver::solvePressureCorrection() {
    std::vector<double> a(numGrids_, 0.0);
    std::vector<double> b(numGrids_, 0.0);
    std::vector<double> c(numGrids_, 0.0);
    std::vector<double> d(numGrids_, 0.0);
    
    // 设置压力边界条件
    pressure_[0] = params_.inletPressure;
    pressure_[numGrids_-1] = params_.outletPressure;
    
    // 构建压力修正方程系数
    for (int i = 1; i < numGrids_ - 1; i++) {
        double dv_dx_g = (V_g_[i+1] - V_g_[i-1]) / (2.0 * dx_);
        double dv_dx_l = (V_l_[i+1] - V_l_[i-1]) / (2.0 * dx_);
        
        // 确保密度不为零
        double rho_g = std::max(1.0, gasDensity_[i]);
        double rho_l = std::max(100.0, oilDensity_[i]);
        
        // 计算系数
        a[i] = -1.0 / (rho_g * dx_ * dx_) - 1.0 / (rho_l * dx_ * dx_);
        c[i] = a[i];
        b[i] = -2.0 * a[i];
        
        // 质量守恒源项
        double mass_imbalance = alpha_g_[i] * dv_dx_g + alpha_l_[i] * dv_dx_l;
        d[i] = -mass_imbalance / dt_;
    }
    
    // 边界条件
    b[0] = 1.0;
    c[0] = 0.0;
    d[0] = 0.0;
    
    a[numGrids_-1] = 0.0;
    b[numGrids_-1] = 1.0;
    d[numGrids_-1] = 0.0;
    
    // 求解三对角矩阵
    solveTridiagonalMatrix(a, b, c, d, pressure_correction_);
    
    // 更新压力和速度
    double omega = 0.2;  // 松弛因子
    
    for (int i = 1; i < numGrids_ - 1; i++) {
        // 限制压力修正量
        double max_correction = 0.05 * pressure_[i];
        pressure_correction_[i] = std::max(-max_correction, 
                                         std::min(max_correction, pressure_correction_[i]));
        
        // 更新压力
        pressure_[i] += omega * pressure_correction_[i];
        
        // 确保压力范围合理
        pressure_[i] = std::max(0.9 * params_.outletPressure, 
                              std::min(1.1 * params_.inletPressure, pressure_[i]));
        
        // 更新速度
        V_g_[i] -= omega * (pressure_[i+1] - pressure_[i-1]) / (2.0 * dx_ * gasDensity_[i]);
        V_l_[i] -= omega * (pressure_[i+1] - pressure_[i-1]) / (2.0 * dx_ * oilDensity_[i]);
        
        // 限制速度范围
        V_g_[i] = std::max(0.5, std::min(10.0, V_g_[i]));
        V_l_[i] = std::max(0.1, std::min(5.0, V_l_[i]));
    }
}

void MultiphaseFlowSolver::solveEnergyEquation() {
    // 简化的能量方程
    for (int i = 1; i < numGrids_ - 1; i++) {
        // 热传导项
        double conduction = (temperature_[i+1] - 2.0*temperature_[i] + temperature_[i-1]) / 
                           (dx_ * dx_);
        
        // 对流项 - 使用上风格式
        double v_mix = alpha_g_[i] * V_g_[i] + alpha_l_[i] * V_l_[i];
        double convection = 0.0;
        if (v_mix > 0) {
            convection = -v_mix * (temperature_[i] - temperature_[i-1]) / dx_;
        } else {
            convection = -v_mix * (temperature_[i+1] - temperature_[i]) / dx_;
        }
        
        // 环境换热项
        double cooling = 0.05 * (params_.ambientTemp - temperature_[i]);
        
        // 更新温度
        double dt_local = std::min(dt_, 0.01);
        temperature_[i] += dt_local * (0.5 * conduction + convection + cooling);
        
        // 限制温度范围
        temperature_[i] = std::max(params_.ambientTemp, 
                                std::min(params_.inletTemp, temperature_[i]));
    }
    
    // 边界条件
    temperature_[0] = params_.inletTemp;
    temperature_[numGrids_-1] = temperature_[numGrids_-2];
}

void MultiphaseFlowSolver::updateProperties() {
    for (int i = 0; i < numGrids_; i++) {
        // 确保压力和温度在合理范围内
        double p = std::max(1000.0, std::min(2.0e7, pressure_[i]));
        double T = std::max(223.15, std::min(373.15, temperature_[i]));
        
        // 更新密度
        gasDensity_[i] = fluid_.getGasDensity(p, T);
        oilDensity_[i] = fluid_.getOilDensity(p, T);
        waterDensity_[i] = fluid_.getWaterDensity(p, T);
        
        // 确保密度不为零
        gasDensity_[i] = std::max(1.0, gasDensity_[i]);
        oilDensity_[i] = std::max(100.0, oilDensity_[i]);
        waterDensity_[i] = std::max(900.0, waterDensity_[i]);
    }
}

void MultiphaseFlowSolver::saveResults(Results& results, double time) {
    for (int i = 0; i < numGrids_; i++) {
        results.pressure[i] = pressure_[i];
        results.temperature[i] = temperature_[i];  // 存储绝对温度
        results.gasHoldup[i] = alpha_g_[i];
        results.oilHoldup[i] = alpha_l_[i];
        results.waterHoldup[i] = alpha_d_[i];
        results.gasVelocity[i] = V_g_[i];
        results.oilVelocity[i] = V_l_[i];
        results.waterVelocity[i] = V_d_[i];
    }

    // 保存详细结果到文件
    std::ofstream outFile("detailed_results.txt", std::ios::app);
    if (outFile.is_open()) {
        for (int i = 0; i < numGrids_; i++) {
            outFile << time << " "
                   << i * dx_ << " "
                   << pressure_[i] << " "
                   << temperature_[i] << " "  // 保存绝对温度
                   << alpha_g_[i] << " "
                   << alpha_l_[i] << " "
                   << alpha_d_[i] << " "
                   << V_g_[i] << " "
                   << V_l_[i] << " "
                   << V_d_[i] << " "
                   << gasDensity_[i] << " "
                   << oilDensity_[i] << " "
                   << waterDensity_[i] << " "
                   << gasViscosity_[i] << " "
                   << oilViscosity_[i] << " "
                   << waterViscosity_[i] << "\n";
        }
        outFile.close();
    }
}

double MultiphaseFlowSolver::calculateFrictionFactor(int i, const std::string& phase) {
    double Re, rho, mu, v;
    if (phase == "gas") {
        rho = gasDensity_[i];
        mu = gasViscosity_[i];
        v = V_g_[i];
    } else {
        rho = oilDensity_[i];
        mu = oilViscosity_[i];
        v = V_l_[i];
    }
    
    double Re_d = rho * std::abs(v) * params_.diameter / mu;
    
    // Blasius摩擦因子关联式
    if (Re_d < 2300) {
        return 64.0 / Re_d;
    } else {
        return 0.316 / std::pow(Re_d, 0.25);
    }
}

void MultiphaseFlowSolver::solveTridiagonalMatrix(
    const std::vector<double>& a,  // 下对角元素
    const std::vector<double>& b,  // 对角元素
    const std::vector<double>& c,  // 上对角元素
    const std::vector<double>& d,  // 右侧向量
    std::vector<double>& x)        // 解向量
{
    int n = a.size();
    std::vector<double> c_prime(n, 0.0);
    std::vector<double> d_prime(n, 0.0);

    // 前向消元
    c_prime[0] = c[0] / b[0];
    d_prime[0] = d[0] / b[0];

    for (int i = 1; i < n; i++) {
        double m = 1.0 / (b[i] - a[i] * c_prime[i-1]);
        c_prime[i] = c[i] * m;
        d_prime[i] = (d[i] - a[i] * d_prime[i-1]) * m;
    }

    // 回代求解
    x[n-1] = d_prime[n-1];
    for (int i = n-2; i >= 0; i--) {
        x[i] = d_prime[i] - c_prime[i] * x[i+1];
    }
}

void MultiphaseFlowSolver::solveMassConservation() {
    // 保存旧的相分数值
    std::vector<double> alpha_g_old = alpha_g_;
    std::vector<double> alpha_l_old = alpha_l_;
    
    double omega = 0.1;  // 松弛因子
    
    for (int i = 1; i < numGrids_ - 1; i++) {
        // 计算通量
        double gas_flux_right = 0.5 * (alpha_g_[i] + alpha_g_[i+1]) * 
                               0.5 * (gasDensity_[i] + gasDensity_[i+1]) * 
                               0.5 * (V_g_[i] + V_g_[i+1]);
        
        double gas_flux_left = 0.5 * (alpha_g_[i-1] + alpha_g_[i]) * 
                              0.5 * (gasDensity_[i-1] + gasDensity_[i]) * 
                              0.5 * (V_g_[i-1] + V_g_[i]);
        
        // 计算相分数变化
        double dalpha_g = -dt_ * (gas_flux_right - gas_flux_left) / 
                         (dx_ * gasDensity_[i]);
        
        // 使用松弛更新气相分数
        alpha_g_[i] = (1.0 - omega) * alpha_g_old[i] + 
                     omega * (alpha_g_old[i] + dalpha_g);
        
        // 限制相分数范围
        alpha_g_[i] = std::max(0.01, std::min(0.99, alpha_g_[i]));
        
        // 更新液相分数
        alpha_l_[i] = 1.0 - alpha_g_[i];
        
        // 检查和修正总和
        double total = alpha_g_[i] + alpha_l_[i];
        if (std::abs(total - 1.0) > 1e-6) {
            alpha_g_[i] /= total;
            alpha_l_[i] /= total;
        }
    }
    
    // 应用边界条件
    alpha_g_[0] = params_.inletGasVolumeFraction;
    alpha_l_[0] = 1.0 - alpha_g_[0];
    
    alpha_g_[numGrids_-1] = alpha_g_[numGrids_-2];
    alpha_l_[numGrids_-1] = alpha_l_[numGrids_-2];
    
    // 确保所有网格点的相分数和为1
    for (int i = 0; i < numGrids_; i++) {
        if (std::isnan(alpha_g_[i]) || std::isnan(alpha_l_[i])) {
            alpha_g_[i] = params_.inletGasVolumeFraction;
            alpha_l_[i] = 1.0 - alpha_g_[i];
        }
        
        if (alpha_g_[i] < 0.0 || alpha_g_[i] > 1.0 ||
            alpha_l_[i] < 0.0 || alpha_l_[i] > 1.0) {
            alpha_g_[i] = params_.inletGasVolumeFraction;
            alpha_l_[i] = 1.0 - alpha_g_[i];
        }
    }
}